import matplotlib.pyplot as plt
import numpy as np
from sklearn import metrics

from readMnist import *

# 数据获取
def get_MNIST_data():
    TRAIN_IMAGES = 'train-images-idx3-ubyte.gz'
    TRAIN_LABELS = 'train-labels-idx1-ubyte.gz'
    TEST_IMAGES = 't10k-images-idx3-ubyte.gz'
    TEST_LABELS = 't10k-labels-idx1-ubyte.gz'
    folder = 'dataset/'

    train_images = extract_images(folder + TRAIN_IMAGES)
    train_labels = extract_labels(folder + TRAIN_LABELS, one_hot=False)

    test_images = extract_images(folder + TEST_IMAGES)
    test_labels = extract_labels(folder + TEST_LABELS, one_hot=False)
    return train_images, train_labels, test_images, test_labels

# k-means初始化类中心
def select_centre(data, k):
    center_list = []
    # 随机
    for _ in range(k):
        ran = np.random.randint(len(data))
        center_list.append(data[ran])
    return center_list

# k-means聚类
def k_means(data, label, center):
    # new_center = []
    label_result = []
    data_result = []
    clusterAssment = np.zeros([len(data), 2])
    label_pre = []
    
    for img, i in zip(data, range(len(data))):
        minDis = 1e9
        minIndex = 0
        for cen, j in zip(center, range(len(center))):
            dis = np.linalg.norm(img - cen)
            if dis < minDis:
                minDis = dis
                minIndex = j
        clusterAssment[i, :] = minIndex, minDis
        label_pre.append(minIndex)

    for j in range(len(center)):
        pointsInCluster = data[clusterAssment[:, 0] == j]
        labelInCluster = label[clusterAssment[:, 0] == j]
        center[j] = np.mean(pointsInCluster, axis=0)
        label_result.append(labelInCluster)
        data_result.append(pointsInCluster)

    return label_pre, label_result, data_result, center

# 训练
def train(data, label, center, iter):
    for _ in range(iter):
        label_pre, label_result, data_result, center = k_means(data, label, center)
        print(metric_acc(label_result, 10))
    return label_pre, label_result, data_result, center

# 测试
def test(data, label, center):
    label_pre, label_result, data_result, center = k_means(data, label, center)
    return label_pre, label_result

# 准确度指标
def metric_acc(label_result, class_num):
    res = [0] * class_num
    acc = 0
    for i in range(class_num):
        res[i] = np.eye(class_num)[label_result[i]].sum(axis=0)
        acc += max(res[i])
    acc /= np.sum(res)
    return acc

# 兰德指数
def metric_ARI(result, pred_result):
    return metrics.adjusted_rand_score(result, pred_result)

if __name__ == '__main__':
    # 加载数据集
    train_data, train_label, test_data, test_label = get_MNIST_data()
    # 初始化类中心
    centers = select_centre(train_data, 10)
    # 训练
    label_pre, label_result, data_result, new_centers = train(train_data, train_label, centers, 10)
    # 测试
    label_pre, label_result = test(test_data, test_label, new_centers)
    # 评价指标
    print('test acc: ' + str(metric_acc(label_result, 10)))
    print('test ARI: ' + str(metric_ARI(test_label, label_pre)))



